const { Ollama } = require("@langchain/ollama");
const { ChatPromptTemplate, HumanMessagePromptTemplate } = require("@langchain/core/prompts");
const { createReactAgent } = require("@langchain/langgraph/prebuilt");
const { HumanMessage } = require("@langchain/core/messages");
const { BufferMemory } = require("langchain/memory");
const { ConversationChain, RetrievalQAChain } = require("langchain/chains");

const hostname = 'http://lbc-ollama.com';

// 初始化llm
const ollama = new Ollama({
    baseUrl: hostname,
    model: "deepseek-r1:1.5b",
    // model: "qwen2.5:1.5b",
    temperature: 0.1,
    numCtx: 4096,
})

// ollama.model = "qwen2.5:1.5b";

// 初始化记忆组件
const memorySaver = new BufferMemory({
    returnMessages: false,
    memoryKey: "chat_history",
    inputKey: "input",
    // outputKey: "output"
});

const prompt = ChatPromptTemplate.fromMessages([
    HumanMessagePromptTemplate.fromTemplate(
        `你是一个非常有帮助的助手。你可以回答任何问题。
        当前对话历史：{chat_history}
        新的输入：{input}
        如果{input} 是一个数字 你将返回该数字加1，返回结果不需要额外输出
        结果: 
        `
    ),
])

// 模板聊天
const chain = new ConversationChain({
    llm: ollama,
    memory: memorySaver,
    prompt: prompt
})

// const qa = RetrievalQAChain


async function saveMemory(params) {
    await memorySaver.saveContext({
        input: params.input,
    }, {
        output: params.output
    });
}

module.exports = {
    ollama,
    chain,
    saveMemory
};